#!/usr/bin/env python
# -*- coding: UTF-8 -*-
'''
@Project ：V2 
@File    ：gas_pipeline_schedule.py
@IDE     ：PyCharm 
@Author  ：郭星
@Date    ：2025/9/23 14:18
我们来重新说一下，有充气和检查两个大步骤，充气部分有两种情况，单独充或者浓度组分相同的气体一起充，
某一瓶所有的组分都充气完成后进行分析，分析部分是可以多个设备串联同时测的，当然也可以单个组分一个一
个的测试，如果多个组分串联设备一起测，就是同时开始测多个组分，可能每个组分的测试时间不一样那就按照最长的那个来，
画甘特图的话，分析部分要稍微明确点，比如有一个气体有两个组分用了设备1设备3，则需要标记出最长的组分时间 在这个时间里不准拆卸加装，
帮我按照上述说得写代码，而且要反馈出来几个参数可以让我输入，比如气瓶数量 充气设备数量 检查设备数量（每个设备只能查一个组分）最大
组分数量（比如如果是4则代表一瓶气可以有一个、两个、三个、最大四个组分），同时甘特图的字体最好能换行显示总是显示不全。按照这个重新写一下
'''
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Gas Mixing -> QC Scheduling (per-bottle; batch in MIX; 1-pass or serial in QC)
------------------------------------------------------------------------------ 
- MIX: same recipe (components+concentrations) can be batched; parallel mixers;
       after MIX of a bottle is done, it can enter QC.
- QC : each device measures exactly ONE component type.
       For a bottle with components {c1,c2,...}:
         * 1-pass: occupy all required devices simultaneously, same start & end,
                   segment length = max(single-device times); NO mid-detach.
         * serial: test devices one-by-one (sum of times).
- Gantt: 1-pass segments printed on all involved devices with same color+label.

Author: you
"""

import random
import numpy as np
from dataclasses import dataclass
from typing import List, Dict, Tuple, Set, Optional

try:
    import matplotlib.pyplot as plt
    HAS_MPL = True
except Exception:
    HAS_MPL = False


# ========================= CONFIG (你可直接改这里) ========================= #

CONFIG = {
    # 规模参数
    "NUM_BOTTLES": 20,                 # 气瓶数量（每瓶视作一个订单）
    "NUM_MIX_MACHINES": 2,             # 充气设备数量
    "NUM_QC_DEVICES": 6,               # 检查设备数量（每台仅测一个组分）
    "MAX_COMPONENTS_PER_BOTTLE": 3,    # 每瓶最大组分数（1~K）

    # 充气参数
    "MIX_BATCH_CAP": 300,              # 同配方批量上限（瓶）
    "MIX_RATE_PER_H": [180.0, 150.0],  # 各充气机速率（瓶/小时），长度>=NUM_MIX_MACHINES
    "MIX_SETUP_SAME": 0.05,            # 同配方/同家族换型时间（小时）
    "MIX_SETUP_DIFF": 0.35,            # 不同配方换型时间（小时）

    # 检查参数（每个组分的“单机检测时间（小时）”）
    # 若 QC 设备数 > 组件种类数，会复用前面的组件类型
    "COMPONENT_TYPES": ["O2", "CO2", "CH4", "CO", "H2", "N2", "C2H6"],
    "COMPONENT_TEST_TIME": {           # 单设备、单瓶、该组分的检测时长
        "O2": 0.18, "CO2": 0.22, "CH4": 0.25, "CO": 0.20, "H2": 0.15, "N2": 0.12, "C2H6": 0.28
    },
    "QC_SETUP_TIME": 0.12,             # 不同“方法签名”（=该瓶在该设备上要测的组分子集，这里就是{component}）切换的校准/换型时间

    # 策略：选择 1-pass 还是 serial
    #   "always_1pass": 所有瓶 1-pass；
    #   "always_serial": 所有瓶串行；
    #   "by_ratio": 按概率 1-pass（见 1PASS_RATIO）；
    #   "by_compare": 若 max(single_times) <= ratio * sum(single_times) 则 1-pass，否则 serial
    "QC_MODE_POLICY": "by_compare",
    "ONEPASS_RATIO": 0.6,              # 用于 "by_compare" 的阈值
    "ONEPASS_RANDOM_RATIO": 0.7,       # 用于 "by_ratio" 的1-pass概率

    # 可视化：甘特文字每行最大字符数（自动换行）
    "LABEL_WRAP_CHARS": 16,
    "RANDOM_SEED": 7
}

# ========================================================================== #


@dataclass(frozen=True)
class Recipe:
    components: Tuple[Tuple[str, float], ...]    # (component, concentration %) sorted

    def key(self) -> Tuple[Tuple[str,float], ...]:
        return self.components


@dataclass
class Bottle:
    id: int
    recipe: Recipe
    due: float = 0.0   # 可选：交期，不影响核心逻辑


@dataclass
class MixTask:
    start: float
    end: float
    machine: int
    label: str


@dataclass
class QCTask:
    start: float
    end: float
    device: int
    label: str
    group_key: Optional[Tuple[int, float, float]] = None  # 用于1-pass分组着色 (bottle_id, st, en)


class RandomData:
    """生成样例数据：瓶、组分、设备->组分映射"""
    def __init__(self, cfg):
        self.cfg = cfg
        random.seed(cfg["RANDOM_SEED"])
        np.random.seed(cfg["RANDOM_SEED"])

    def gen(self):
        N = self.cfg["NUM_BOTTLES"]
        D = self.cfg["NUM_QC_DEVICES"]
        comps_all = self.cfg["COMPONENT_TYPES"]
        Kmax = self.cfg["MAX_COMPONENTS_PER_BOTTLE"]

        # 给每台QC设备指定一种“唯一组件类型”（可能会循环复用前面类型）
        dev_to_comp: List[str] = []
        for i in range(D):
            dev_to_comp.append(comps_all[i % len(comps_all)])

        # 生成瓶的配方（确保都能被设备覆盖）
        bottles: List[Bottle] = []
        for i in range(N):
            k = np.random.randint(1, Kmax+1)  # 本瓶组分数
            # 从设备可测的集合抽取，以保证可测
            pickable = list(dict.fromkeys(dev_to_comp))  # 保持顺序去重
            comps = random.sample(pickable, k)
            # 生成浓度（演示：每个 0.1~5%）
            comps_t = tuple(sorted(((c, round(np.random.uniform(0.1, 5.0), 3)) for c in comps), key=lambda x:x[0]))
            bottles.append(Bottle(i, Recipe(comps_t)))
        return bottles, dev_to_comp


# ============================ 核心：调度器 ============================ #

class Scheduler:
    def __init__(self, cfg, bottles: List[Bottle], dev_to_comp: List[str]):
        self.cfg = cfg
        self.bottles = bottles
        self.N = len(bottles)
        self.dev_to_comp = dev_to_comp
        self.D_mix = cfg["NUM_MIX_MACHINES"]
        self.D_qc  = cfg["NUM_QC_DEVICES"]

        # 若充气速率列表短于设备数，循环使用
        rates = cfg["MIX_RATE_PER_H"]
        if len(rates) < self.D_mix:
            mul = (self.D_mix + len(rates) - 1) // len(rates)
            rates = (rates * mul)[:self.D_mix]
        self.mix_rates = rates

    # ---------- 充气阶段：按配方合批并排产 ----------
    def schedule_mixing(self) -> Tuple[Dict[int, float], List[MixTask]]:
        """
        批处理规则：
          - 同 recipe.key() 的瓶可合批，批量<=MIX_BATCH_CAP。
          - 派工：哪个混气机最早能完工就派哪台。
          - 工时 = 批内瓶数 / 该机速率；换型：相同配方=MIX_SETUP_SAME, 不同=MIX_SETUP_DIFF。
        返回：每瓶 mix_end 时间；MIX 甘特列表。
        """
        cap = self.cfg["MIX_BATCH_CAP"]
        setup_same = self.cfg["MIX_SETUP_SAME"]
        setup_diff = self.cfg["MIX_SETUP_DIFF"]

        # 按配方聚类
        key2idx: Dict[Tuple, List[int]] = {}
        for b in self.bottles:
            key2idx.setdefault(b.recipe.key(), []).append(b.id)

        # 每个簇切成若干批（<=cap）
        batches: List[List[int]] = []
        for key, ids in key2idx.items():
            cur=[]
            for bid in ids:
                cur.append(bid)
                if len(cur) >= cap:
                    batches.append(cur); cur=[]
            if cur: batches.append(cur)

        # 这里只做一个稳定顺序：按簇大小从大到小
        batches.sort(key=lambda bl: -len(bl))

        # 每台机的 ready 与上一配方key
        ready = [0.0]*self.D_mix
        last_key = [None]*self.D_mix
        mix_end: Dict[int, float] = {}
        tasks: List[MixTask] = []

        for bl in batches:
            # 批的“代表配方key”
            rep_key = self.bottles[bl[0]].recipe.key()
            # 选机：完工时间最早
            best = None
            for m in range(self.D_mix):
                setup = setup_same if last_key[m] == rep_key else (0.0 if last_key[m] is None else setup_diff)
                st = ready[m] + setup
                pt = len(bl) / self.mix_rates[m]
                en = st + pt
                if (best is None) or (en < best[0]):
                    best = (en, m, st, pt, setup)
            en, m, st, pt, setup = best
            # 记录
            tasks.append(MixTask(st, en, m, f"Batch({','.join('S'+str(i) for i in bl)})"))
            ready[m] = en; last_key[m] = rep_key
            for bid in bl:
                mix_end[bid] = en

        return mix_end, tasks

    # ---------- QC 单机单瓶时长 ----------
    def single_device_time(self, comp: str) -> float:
        tmap = self.cfg["COMPONENT_TEST_TIME"]
        return float(tmap.get(comp, 0.2))

    # ---------- 决策：此瓶 1-pass 还是 serial ----------
    def choose_qc_mode(self, comps: List[str]) -> str:
        policy = self.cfg["QC_MODE_POLICY"]
        single_times = [self.single_device_time(c) for c in comps]
        mx, sm = max(single_times), sum(single_times)

        if policy == "always_1pass":
            return "1pass"
        if policy == "always_serial":
            return "serial"
        if policy == "by_ratio":
            return "1pass" if random.random() < self.cfg["ONEPASS_RANDOM_RATIO"] else "serial"
        # by_compare
        return "1pass" if mx <= self.cfg["ONEPASS_RATIO"] * sm else "serial"

    # ---------- QC 调度 ----------
    def schedule_qc(self, mix_end: Dict[int, float]) -> Tuple[float, List[QCTask]]:
        """
        规则：
          - 每台QC设备仅测一个组分（dev_to_comp 已给定）。
          - 一瓶进入QC的最早时间 = mix_end[bid]
          - 1-pass：找到该瓶所需的全部设备，统一开始时刻 st = max(各设备 ready+setup, mix_end)；
                    段长 = max(各设备的单机检测时长)；所有设备上都画同一段 [st,en]。
          - serial：按设备ID升序逐台；每步 st = max(上一台完工, 该设备 ready+setup)；段长=该设备单机时长。
          - setup：若该设备上一单的“方法签名”（此处=该瓶该设备上要测的组分名称）不同，则加 QC_SETUP_TIME。
        """
        D = self.D_qc
        ready = [0.0]*D
        last_sig: List[Optional[str]] = [None]*D
        setup_t = self.cfg["QC_SETUP_TIME"]
        tasks: List[QCTask] = []

        # 瓶按 MIX 完工早的先做（也可改为你的序策略）
        order_seq = sorted(self.bottles, key=lambda b: mix_end[b.id])

        for b in order_seq:
            bid = b.id
            need_comps = [c for c,_ in b.recipe.components]

            # 找到每个需要的组件对应的设备列表（可能多个设备同类型？这里每类只一台）
            comp_to_dev: Dict[str, List[int]] = {}
            for d, ctype in enumerate(self.dev_to_comp):
                comp_to_dev.setdefault(ctype, []).append(d)
            # 对每个需要的组分，取一台可测的设备（简单选“最早可用”的那台）
            devs: List[int] = []
            for comp in need_comps:
                cand = comp_to_dev.get(comp, [])
                if not cand:
                    raise RuntimeError(f"组件 {comp} 无可用设备！")
                # 选 ready+setup 最早的那台
                best_d = None; best_ready = None
                for d in cand:
                    sig = comp         # 方法签名简化为该组件名
                    stp = 0.0 if (last_sig[d] is None or last_sig[d] == sig) else setup_t
                    rdy = ready[d] + stp
                    if (best_ready is None) or (rdy < best_ready):
                        best_ready = rdy; best_d = d
                devs.append(best_d)

            # 选择模式
            mode = self.choose_qc_mode(need_comps)

            if mode == "1pass":
                # 统一开始：所有设备就绪与 mix_end 取最大值
                earliest = mix_end[bid]
                per_pt = []
                stp_sum = 0.0
                for d, comp in zip(devs, need_comps):
                    sig = comp
                    stp = 0.0 if (last_sig[d] is None or last_sig[d] == sig) else setup_t
                    stcand = ready[d] + stp
                    earliest = max(earliest, stcand)
                    pt = self.single_device_time(comp)
                    per_pt.append((d, sig, stp, pt))
                seg_len = max(pt for (_,_,_,pt) in per_pt)
                st = earliest; en = st + seg_len
                group_key = (bid, round(st,6), round(en,6))
                label = f"S{bid} 1-p [{'+'.join('E'+str(d+1) for (d,_,_,_) in per_pt)}]\nmax={seg_len:.2f}h"
                for d, sig, stp, pt in per_pt:
                    tasks.append(QCTask(st, en, d, label, group_key))
                    ready[d] = en
                    last_sig[d] = sig
                    stp_sum += stp
                # 这里不单独累计/输出setup时间；如需可返回
            else:
                # serial 逐台（按设备ID升序）
                cur = mix_end[bid]
                for d, comp in sorted(zip(devs, need_comps), key=lambda x:x[0]):
                    sig = comp
                    stp = 0.0 if (last_sig[d] is None or last_sig[d] == sig) else setup_t
                    st = max(cur, ready[d] + stp)
                    pt = self.single_device_time(comp)
                    en = st + pt
                    tasks.append(QCTask(st, en, d, f"S{bid} s E{d+1} ({comp})", None))
                    ready[d] = en
                    last_sig[d] = sig
                    cur = en

        Cmax = 0.0
        if tasks:
            Cmax = max(t.end for t in tasks)
        return Cmax, tasks


# ============================ 可视化（甘特） ============================ #

def wrap_text(s: str, width: int) -> str:
    """简单按字符数换行（不掰单词）"""
    lines=[]
    cur=''
    for ch in s:
        cur += ch
        if len(cur) >= width and ch == ' ':
            lines.append(cur.rstrip()); cur=''
    if cur: lines.append(cur)
    return '\n'.join(lines)

def plot_gantt(mix_tasks: List[MixTask], qc_tasks: List[QCTask],
               cfg, title="Mix & QC Gantt (1-pass segments same on all devices)"):
    if not HAS_MPL:
        print("(matplotlib not installed; skip plotting)"); return

    # 行：MIX 设备 + QC 设备
    rows: List[Tuple[str, int]] = []
    for m in range(cfg["NUM_MIX_MACHINES"]):
        rows.append(("MIX", m))
    for d in range(cfg["NUM_QC_DEVICES"]):
        rows.append(("QC", d))

    row_index = {("MIX",m): i for i,m in enumerate(range(cfg["NUM_MIX_MACHINES"]))}
    base_qc = len(row_index)
    for d in range(cfg["NUM_QC_DEVICES"]):
        row_index[("QC",d)] = base_qc + d

    # 颜色：对 1-pass 的 group_key 统一颜色
    from itertools import cycle
    palette = cycle(plt.rcParams['axes.prop_cycle'].by_key()['color'])
    group_color: Dict[Tuple[int,float,float], str] = {}
    for t in qc_tasks:
        if t.group_key and t.group_key not in group_color:
            group_color[t.group_key] = next(palette)

    fig_h = max(3.5, 0.6 * len(rows) + 2)
    fig, ax = plt.subplots(figsize=(12, fig_h))

    # 画 MIX
    for mt in mix_tasks:
        y = row_index[("MIX", mt.machine)]
        ax.barh(y, mt.end-mt.start, left=mt.start, height=0.8)
        ax.text((mt.start+mt.end)/2, y,
                wrap_text(mt.label, cfg["LABEL_WRAP_CHARS"]),
                ha='center', va='center', color='black', fontsize=8)

    # 画 QC
    for qt in qc_tasks:
        y = row_index[("QC", qt.device)]
        color = group_color.get(qt.group_key, None)
        ax.barh(y, qt.end-qt.start, left=qt.start, height=0.8,
                color=color)
        ax.text((qt.start+qt.end)/2, y,
                wrap_text(qt.label, cfg["LABEL_WRAP_CHARS"]),
                ha='center', va='center',
                color=('black' if color else 'black'), fontsize=6)

    # y 轴标签
    yticks = []
    ylabels = []
    for m in range(cfg["NUM_MIX_MACHINES"]):
        yticks.append(row_index[("MIX", m)])
        ylabels.append(f"MIX-M{m+1}")
    for d in range(cfg["NUM_QC_DEVICES"]):
        yticks.append(row_index[("QC", d)])
        ylabels.append(f"QC-E{d+1} ({CONFIG['COMPONENT_TYPES'][d % len(CONFIG['COMPONENT_TYPES'])]})")

    ax.set_yticks(yticks); ax.set_yticklabels(ylabels)
    ax.set_xlabel("Time (h)"); ax.set_title(title)
    ax.grid(axis='x', alpha=0.3)
    plt.tight_layout()
    plt.show()


# ============================ 主流程 ============================ #

def main():
    # 1) 生成数据
    data = RandomData(CONFIG)
    bottles, dev_to_comp = data.gen()

    # 2) 排程
    sch = Scheduler(CONFIG, bottles, dev_to_comp)
    mix_end, mix_tasks = sch.schedule_mixing()
    Cmax, qc_tasks = sch.schedule_qc(mix_end)

    # 3) 输出与画图
    print(f"Total bottles: {CONFIG['NUM_BOTTLES']}")
    print(f"Mix machines: {CONFIG['NUM_MIX_MACHINES']}, QC devices: {CONFIG['NUM_QC_DEVICES']}")
    print(f"Cmax (QC finish): {Cmax:.3f} h")

    if HAS_MPL:
        plot_gantt(mix_tasks, qc_tasks, CONFIG)
    else:
        print("(Install matplotlib to see the Gantt chart.)")

if __name__ == "__main__":
    main()
